[DCP] Run fi_a2a on FlashInfer's fused all-to-all + LSE reduce by thanhhao98 · Pull Request #40539 · sgl-project/sglang

thanhhao98 · 2026-09-21T06:49:23Z

Motivation

--dcp-comm-backend fi_a2a (the Blackwell default for MLA decode context parallelism) runs three kernels per MLA layer per
decode step:

dcp_pack_a2a_send packs the attention output;
FlashInfer's MNNVL helixAllToAllKernel exchanges it;
the Triton _dcp_lse_combine_kernel merges it.

FlashInfer's decode_cp_a2a_lse_reduce does the exchange and the LSE-weighted merge in one kernel:

feat(comm): add fused decode CP A2A + LSE reduce flashinfer-ai/flashinfer#4929 added it;
perf(comm): LL128 protocol for decode_cp_a2a_lse_reduce, on any torch symmetric-memory backend flashinfer-ai/flashinfer#5726 rewrites its protocol (LL128: no fence or grid barrier on the critical path) and makes it run
on whichever torch symmetric-memory backend the process uses.

This PR makes fi_a2a use that op, and only it.

The kernel reads the attention output in place (a strided view), so the pack kernel is gone too: three kernels per layer
become one.
The workspace lives on torch's default (CUDA) symmetric-memory backend, which it shares with CustomAllReduceV2. Nothing
else in the process changes: no backend switch, custom all-reduce v2 stays on, no NCCL_CUMEM_ENABLE change.

Dependency. This needs a flashinfer-python release that contains flashinfer-ai/flashinfer#5726. sglang currently pins
0.7.0.post1, which has no fused op; the release candidates so far (0.7.1rc1–rc5, 0.7.2rc1) have #4929's version, which
forces the NCCL backend. This PR stays draft until the FlashInfer bump lands. Until then:

an unset --dcp-comm-backend resolves to a2a and logs why;
an explicit fi_a2a is rejected at argument resolution;
a FlashInfer whose op still switches the backend fails at startup with a message naming the fix.

Modifications

fi_a2a = the fused kernel, exclusively. The MNNVL exchange + Triton merge path is removed.
Startup, before CUDA-graph capture (init_fi_a2a_workspace):
- creates two workspaces, one for graph capture and replay and one for eager calls, because a workspace serves one ordered
  CUDA stream;
- sizes them for the widest decode CUDA-graph batch (--cuda-graph-max-bs-decode, times the tokens each request verifies
  under speculative decoding), also with --disable-decode-cuda-graph;
- logs DCP fi_a2a: fused all-to-all + LSE reduce ready (backend=..., capacity ..., 2 workspaces x ... MiB) once.
Per call (fi_a2a_lse_reduce, a registered custom op, so torch.compile sees one opaque collective):
- views the [B, H, D] attention output as [B, H/cp, cp, D] and its LSE as [B, H/cp, cp] (no copy); head h goes to
  peer h // (H/cp), exactly as the pack kernel sent it;
- picks lse_mode from the attention backend;
- runs eager batches wider than the workspace as several calls, in the same order on every DCP rank.
Resolution.
- fi_a2a_platform_blocker: Blackwell, the DCP group inside one NVLink domain (GB200/GB300 fabric or one node), and a
  FlashInfer that provides the op. It never imports FlashInfer.
- fi_a2a_blocker adds an MLA model, no --enable-pdmux, and no PP DeepGEMM warmup stream at --pp-size > 1.
- Unset resolves to fi_a2a only where it can run, else a2a, and logs the reason. An explicit fi_a2a names the blocker.
Docs. dcp.mdx, server_arguments.mdx, the --dcp-comm-backend help, and the Kimi-K3 cookbook note.

Accuracy Tests

New GPU test test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py, at 2 / 4 / 8 × B200: 10 passed,
0 skipped at each size. It checks:
- equivalence with the all-gather + Triton merge, base-2 and base-e, at 16 heads × 3 tokens, 16 × 7 (several calls per
  batch) and 2 × 49;
- FlashMLA's [B, H, 1] LSE;
- CUDA-graph replays interleaved with eager calls;
- a third CUDA stream is rejected before the collective;
- custom all-reduce v2 stays exact around fused calls on the same symmetric memory.
2 × 4 GB300, the DCP group across two trays of one NVL72: the test passes on all 8 ranks (9 passed; the all-reduce case
is skipped because sglang does not build custom all-reduce v2 for that group). DeepSeek-V2-Lite TP 8 = DCP 8: GSM8K 0.640
for both fi_a2a and a2a, and TPOT at concurrency 1 is 2.3% lower with fi_a2a.
GSM8K (5-shot, 50 questions, --parallel 1) on every serving arm below: main 0.80 and this PR 0.76 in every run
(2 of 50 questions; greedy decoding repeats exactly). The full question set is below.
Greedy canary: four prompts give token-identical outputs on main and this PR: 8 of 8 cross-path run pairs.
Coverage of the narrow-graph and eager paths (GSM8K, 1319 questions, --parallel 32):
- --cuda-graph-max-bs-decode 4: this PR 0.772, main 0.766;
- this PR without decode CUDA graphs (eager decode, several calls per batch): 0.767.

Speed Tests and Profiling

8× B200, DeepSeek-V2-236B BF16, TP 8 = DCP 8, FlashInfer MLA, 256K context, image
lmsysorg/sglang:nightly-dev-cu13-20261006-f3fe7534
(torch 2.14.1, sglang-kernel 0.4.9). FlashInfer 0.7.0.post1 with the op from flashinfer-ai/flashinfer#5726 laid over it, the
same for every arm. Default settings in every arm (custom all-reduce v2 on).

Per MLA layer (microbenchmark, 16 heads per rank, head_dim 512, 60 calls per CUDA graph, max over ranks of each rank's
median, µs):

batch	main: pack + MNNVL + combine	this PR	Δ
1	8.08	6.55	−18.9%
2	8.21	6.85	−16.6%
4	8.82	7.27	−17.6%
8	10.68	8.20	−23.2%
16	15.05	10.50	−30.2%
32	20.60	13.97	−32.2%
64	32.38	22.38	−30.9%
128	54.47	39.23	−28.0%

L2 hot, mean of two repeats in one job (palindromic order); with L2 flushed before every layer, 9.51 → 7.25 µs at batch 1.

Main's three kernels, timed one at a time (another job, same node type):
- pack 1.47 + all-to-all 3.86 + combine 1.58 = 6.91 µs at batch 1;
- the same three chained cost 8.09 µs in that job. The difference is the gaps between the dependent kernels (nsys in
  serving: 0.26 + 0.60 µs).
FlashInfer's current op ([Bug] sglang frontend has an incorrect dependency on torch #4929), same job: 27.4 µs at batch 1 and 37.5 µs at batch 32, on contiguous input.

Serving (median TPOT, ms, ISL 260992 / OSL 1024). Each job runs main, this PR, this PR, main (or the reverse) on one
node:

job (order)	concurrency 1: main	this PR	Δ	concurrency 2 Δ	median ITL Δ, concurrency 1 / 2 / 4
1 (main, PR, PR, main)	13.271	13.144	−0.96%	−0.40%	−0.93 / −0.66 / −1.03%
2 (PR, main, main, PR)	13.227	13.120	−0.81%	−0.63%	−0.88 / −0.90 / −0.79%
3 (main, PR, PR, main)	13.246	13.127	−0.90%	−0.66%	−0.90 / −0.95 / −0.81%
4 (PR, main, main, PR)	13.244	13.137	−0.81%	−1.43%	−0.84 / −0.88 / −0.87%
mean			−0.87% (4/4)	−0.78%	−0.89 / −0.85 / −0.88%

Decode-only step time at 32K context (one_batch_server), batch 8 / 16 / 32: −1.19 / −1.14 / −1.52%, lower in all 4
jobs.
Concurrency-4 TPOT is noisy: only 9 requests finish per run, and TPOT includes stalls behind other requests'
261K-token prefills.
- Six of the eight runs of this PR are 0.2–0.5% below main.
- Two read 46.9 and 58.2 ms (main: 44.5–44.7) while their median ITL, 21.56 and 21.51 ms, matches the other runs of
  this PR.
- Median ITL at concurrency 4 is −0.88% (−0.79 to −1.03% per job).
Control (a2a on main vs this PR, untouched by this change): +0.05% TPOT at concurrency 1.
vs a2a (why the resolver keeps preferring fi_a2a): TPOT at concurrency 1 is 5.1% lower with fi_a2a (13.16 vs
13.86 ms).
Why the FlashInfer version matters. FlashInfer's current op ([Bug] sglang frontend has an incorrect dependency on torch #4929, in the 0.7.1 / 0.7.2 release candidates) cannot
start next to custom all-reduce v2. With v1 all-reduce in every arm, TPOT at concurrency 1 is:
- main 13.678 ms;
- the PR's path on [Bug] sglang frontend has an incorrect dependency on torch #4929's op 15.658 ms (+14.5%), with the inputs copied contiguous;
- this PR on flashinfer#5726 13.570 ms (−0.79%).
nsys puts the DCP time per decode step at 568, 2534 and 445 µs respectively.

nsys (20 decode steps at concurrency 1; two captures in opposite orders; kernels per decode step on each of the 8 GPUs):

	main	this PR
`_dcp_pack_a2a_send_kernel`	60 (median 1.50 µs)	0
`helixAllToAllKernel`	60 (5.44 µs)	0
`_dcp_lse_combine_kernel`	60 (1.70 µs)	0
`flashinfer::comm::dcp::FusedKernel` (base-2)	0	60 (7.26 µs)
custom all-reduce v2 kernels	121	121
critical path per layer (last rank's attention end → last rank's reduce end)	9.41 µs	7.54 µs
× 60 layers, per decode step	566 µs of 13.22 ms (4.3%)	453 µs of 13.07 ms (3.5%)

The three measurements agree. Per decode step at batch 1:
- the microbenchmark predicts −93 µs (L2 hot) to −141 µs (L2 flushed);
- nsys measures −113 µs (−112 µs in the reverse-order capture);
- serving measures −115 µs (concurrency-1 TPOT).
- At batch 8 / 16 / 32 the decode-only step shrinks by 177 / 230 / 445 µs, against a predicted 146–194 / 266–272 /
  399–410 µs.
Memory. Two workspaces of 9.4 MiB each at --cuda-graph-max-bs-decode 32. They grow with the widest captured decode
batch times the local heads, and are allocated after KV-cache sizing, so max_total_num_tokens is unchanged (11,465,208 in
both arms). dcp.mdx documents the size.

Coordination

[DSA] Fix decode context parallelism for RoPE sparse attention on Blackwell #42536 (DSA DCP). Its fi_a2a memory term assumes the MNNVL workspace. The fused workspace is sized by
decode_cp_a2a_lse_reduce_workspace_size(...) (two of them), and wider batches are chunked.
[ROCm] Support decode context parallel (DCP) for GLM-5 / DeepSeek-V3.2 #42618. Touches only the a2a branch of dcp_a2a_lse_reduce.
[DCP][Spec] EAGLE/EAGLE3 support for decode context parallelism #31785. Overlaps only textually (dcp.mdx, test_server_args.py).

Checklist

Format your code according to the Code Formatting with Pre-Commit.
Add unit tests according to the Running and Adding Unit Tests.
Update documentation according to Write Documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark and Profiling.

CI States

Latest PR Test (Base): ❌ Run #37612939714
Latest PR Test (Extra): ❌ Run #37612939327
Latest PR Test (AMD ROCm 10): ❌ Run #37612939559

…E reduce The existing a2a and fi_a2a backends exchange head partials and then merge them with a separate Triton LSE kernel, so every MLA layer round-trips the partials through global memory between the two kernels. FlashInfer's fused decode-CP op does the exchange and the LSE-weighted reduce inside one cooperative kernel over NCCL device-side LSA symmetric memory, removing that round-trip and one launch per layer per decode step. Expose it as a fourth --dcp-comm-backend value rather than changing fi_a2a in place, so the validated MNNVL path is untouched and the two can be compared directly. Resolution never selects it on its own; it is explicit-only. The fused op allows one workspace per ordered CUDA stream, so the workspaces are allocated upfront (rendezvous is collective and not graph-capturable) and handed out per stream: one for eager decode, one for the global graph-capture stream. Requires a FlashInfer build carrying the op (flashinfer-ai/flashinfer#4929), NCCL >= 2.29 with device-communicator support, and a DCP group inside one NVLink domain. fp16/bf16 only -- unlike a2a, which ships partials as opaque bytes, the fused kernel has no fp8 path. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…A2A + LSE reduce" This reverts commit 852d660.

fi_a2a is moving to FlashInfer's decode_cp_a2a_lse_reduce, which needs torch >= 2.14, a FlashInfer build that ships the op, an MLA model and a DCP group inside one NVLink domain. It has no fallback kernel, so an unset --dcp-comm-backend must only resolve to fi_a2a where all of that holds, and otherwise to a2a. fi_a2a_platform_blocker reports the host-side reason without importing FlashInfer or creating a CUDA context; fi_a2a_blocker adds the configuration checks (MLA, PDMUX, torch symm-mem all-reduce, the PP DeepGEMM warmup stream, an exported NCCL_CUMEM_ENABLE=0). An explicit fi_a2a names the blocker, and otherwise defaults NCCL_CUMEM_ENABLE=1 because the fused workspace is an NCCL symmetric window.

fi_a2a exchanged DCP head partials with FlashInfer's MNNVL all-to-all and then merged them with a separate Triton LSE kernel, so every MLA layer round-tripped the partials through HBM between two launches. decode_cp_a2a_lse_reduce does the exchange and the LSE-weighted merge in one cooperative kernel over NCCL LSA symmetric memory; fi_a2a now uses it exclusively. Before graph capture, every DCP rank agrees over gloo that the op can run, then creates two workspaces (one per ordered CUDA stream: graph capture and serving) and checks one probe call per LSE base against an all-gather + Triton reference. That probe also makes the group's first, collective device-comm setup eager. Batches wider than the captured graphs run as row chunks on the same stream and workspace. The op is wrapped as a custom op so Dynamo does not trace the stream-to-workspace lookup. The workspace comes from torch symmetric memory on the NCCL backend, a process-global choice frozen by the first symmetric allocation, so bootstrap selects it before the groups are built. CustomAllReduceV2 cannot allocate from that backend (it passes a group name and rendezvouses on gloo), so it yields to CustomAllreduce there.

Describe what fi_a2a now requires (Blackwell, one NVLink domain, an MLA model, torch >= 2.14, a FlashInfer build with decode_cp_a2a_lse_reduce), how an unset --dcp-comm-backend resolves, and the process-wide settings it changes: NCCL symmetric-memory backend with custom all-reduce v1, NCCL_CUMEM_ENABLE=1, and PDMUX / torch symm-mem all-reduce rejected.

dongha-yoon · 2026-09-30T23:18:27Z

Hi, @thanhhao98 flashinfer-ai/flashinfer#4929 has been merged.

FlashInfer's fused DCP reduce allocates its workspace on whichever torch symmetric-memory backend is active. Stop selecting NCCL at bootstrap, and drop the custom all-reduce v1 fallback, the NCCL_CUMEM_ENABLE default and the cuMem, torch 2.14, NCCL 2.29 and --enable-torch-symm-mem checks that only that backend needed. No per-rank environment check is left to agree on, so remove the startup probe and its gloo vote. Keep one workspace for graphs captured on the capture stream and one for eager calls, chosen by the current stream, and re-raise a failed workspace creation with what to change.

The fused reduce accepts strided [B, H/cp, cp, D] and [B, H/cp, cp] views, so view the attention output and LSE instead of packing them. Head h still goes to peer h // (H/cp). Inputs with a non-unit inner stride, or rows off an 8-byte boundary, are copied first; wide eager batches still run as token slices. Drop the pack test for the layout only fi_a2a used.

fi_a2a needs Blackwell, an MLA model, a DCP group inside one NVLink domain and a flashinfer-python with the fused decode_cp_a2a_lse_reduce; it no longer changes the symmetric-memory backend, custom all-reduce or NCCL settings, so drop those notes. Describe the two workspaces, the chunked eager batches and the configurations that resolve to a2a.

The DCP docs now give the size of the two fused-reduce workspaces: they follow --cuda-graph-max-bs-decode, also with decode CUDA graphs disabled, and come out of the memory left after KV-cache sizing. The Kimi K3 notes name the FlashInfer op the default backend needs. Drop the copy for strided or misaligned partials: the only caller passes contiguous ones, and FlashInfer checks its inputs. Comments that still described fi_a2a as packing and running the Triton combine now describe a2a only. Tests: keep the bf16 reference cases that cover the head mapping and chunking, and drop the cases that only exercised FlashInfer or the removed copy. Fail rather than skip when custom all-reduce v2 should be running. Add the PP case that keeps fi_a2a without the DeepGEMM warmup.

…rves the group Cross-node groups use custom all-reduce v2 only when they form one NVLink clique on a VMM-backed allocator, so the coexistence case now asks the same predicate the dispatcher uses and skips otherwise.

dongha-yoon · 2026-10-07T17:29:07Z

Hi @thanhhao98, is there anything I can help with to move this PR forward?

github-actions Bot added the amd label Sep 21, 2026

thanhhao98 force-pushed the htphan/dcp-fi-a2a-fused branch from 9ef5f8f to 852d660 Compare September 21, 2026 08:47

kpham-sgl self-assigned this Sep 23, 2026

Hao Phan added 5 commits September 30, 2026 01:02

Revert "[DCP] Add fi_a2a_fused comm backend using FlashInfer's fused …

dbeafbf

…A2A + LSE reduce" This reverts commit 852d660.

Merge remote-tracking branch 'origin/main' into htphan/dcp-fi-a2a-fused

89d85dc

Hao Phan added 7 commits October 6, 2026 15:59

Merge remote-tracking branch 'origin/main' into htphan/dcp-fi-a2a-fused

d6f9c5b

Merge remote-tracking branch 'origin/main' into htphan/dcp-fi-a2a-fused

f029f84

github-actions Bot added documentation Improvements or additions to documentation jit-kernel labels Oct 7, 2026

thanhhao98 changed the title ~~[DCP] Add fi_a2a_fused comm backend using FlashInfer's fused A2A + LSE reduce~~ [DCP] Run fi_a2a on FlashInfer's fused all-to-all + LSE reduce Oct 7, 2026

Navigation Menu

Repository navigation

[DCP] Run fi_a2a on FlashInfer's fused all-to-all + LSE reduce - #40539

[DCP] Run fi_a2a on FlashInfer's fused all-to-all + LSE reduce#40539
thanhhao98 wants to merge 13 commits into
sgl-project:mainfrom
thanhhao98:htphan/dcp-fi-a2a-fused

thanhhao98 commented Sep 21, 2026 •

edited by github-actions Bot

Loading

Uh oh!

dongha-yoon commented Sep 30, 2026

Uh oh!

dongha-yoon commented Oct 7, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

thanhhao98 commented Sep 21, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Coordination

Checklist

CI States

Uh oh!

dongha-yoon commented Sep 30, 2026

Uh oh!

dongha-yoon commented Oct 7, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

thanhhao98 commented Sep 21, 2026 •

edited by github-actions Bot

Loading